// "Copyright [2021] <Copyright Shengkai Lin>"

// This file is some basic class and useful functions

#ifndef CLUSTER_BASIC_BASIC_H_
#define CLUSTER_BASIC_BASIC_H_

#include <omnetpp.h>
#include <vector>
#include <unordered_map>
#include <memory>
#include <utility>
#include "cluster/basic/topology.h"

namespace ai_sim {

using namespace std;

typedef vector<pair<int, int>> IntPairVector;
typedef map<int, int> IntMap;

// Used in resource_scheduler
enum GpuScheduleAlgo {
    Sequential,
    Worstcase,
    Bestcase,
    Randomly,
    Manually
};

enum TaskType {
    RingAllreduce,
    DoubleBinaryTree,
    Butterfly,
    HD
};

enum ClusterType {
    OpticalCluster,
    ElectricalCluster,
    HybridCluster,
    IbCluster
};

enum SpecialRequireofFlow {
    NoRequirement,
    NeedControllerCooperate,  // Need to inform the controller before send
    NeedCompute  // Need to wait to compute
};

// Used in switches
struct RouteInfo {
    int inPort;
    int outPort;
    int srcIdx;
    int destIdx;
};

typedef vector<ai_sim::RouteInfo> RouteVec;

class Task {
 public:
    int taskId;
    TaskType tasktype;
    int numGpus;
    int numBatches;
    int modelsize;
    double compute_time;

 public:
    Task() {
        // Gen a empty task
        srand(0);
        taskId = rand();
        tasktype = RingAllreduce;
        numGpus = 2;
        numBatches = 1;
        modelsize = 0;
        compute_time = 0;
    }
    Task(int taskId_, TaskType tasktype_, int numGpus_,
        int numBatches_, int modelsize_, double compute_time_) {
        taskId = taskId_;
        tasktype = tasktype_;
        numGpus = numGpus_;
        numBatches = numBatches_;
        modelsize = modelsize_;
        compute_time = compute_time_;
    }
};

class Flow {
 public:
    int flowId;
    Task the_task;
    double data_size;
    int srcIdx;
    int destIdx;
    int outportId;
    int numUpstreamflows;
    vector<int> downstream_flows;
    map<int, RouteVec> dependRoute;
    // the bool represents that if the require is completed
    unordered_map<SpecialRequireofFlow, bool> requires;
    int offset = 0;
    bool isLastPacket = false;

 public:
    Flow() {
        srand(0);
        flowId = rand();
        data_size = 0;
        srcIdx = -1;
        destIdx = -1;
        outportId = -1;
        numUpstreamflows = 0;
        unordered_map<SpecialRequireofFlow, bool> empty_require;
        empty_require[NeedCompute] = false;
        empty_require[NeedControllerCooperate] = false;
        requires = empty_require;
        offset = 0;
        isLastPacket = false;
    }
    Flow(double data_size_, int srcIdx_, int destIdx_, int outportId_) {
        srand(0);
        flowId = rand();
        data_size = data_size_;
        srcIdx = srcIdx_;
        destIdx = destIdx_;
        outportId = outportId_;
        numUpstreamflows = 0;
        unordered_map<SpecialRequireofFlow, bool> empty_require;
        empty_require[NeedCompute] = false;
        empty_require[NeedControllerCooperate] = false;
        requires = empty_require;
        offset = 0;
        isLastPacket = false;
    }
    Flow(int flowId_,
        Task the_task_,
        double data_size_,
        int srcIdx_,
        int destIdx_,
        int outportId_,
        int numUpstreamflows_,
        vector<int> downstream_flows_,
        map<int, RouteVec> dependRoute_,
        // the bool represents that if the require is completed
        unordered_map<SpecialRequireofFlow, bool> requires_) :
    flowId(flowId_), the_task(the_task_), data_size(data_size_),
    srcIdx(srcIdx_), destIdx(destIdx_), outportId(outportId_),
    numUpstreamflows(numUpstreamflows_), downstream_flows(downstream_flows_),
    dependRoute(dependRoute_), requires(requires_){}
};

typedef std::shared_ptr<ai_sim::Flow> SharedFlowPtr;

typedef vector<Flow> FlowVector;
typedef vector<int> IntVector;

enum ProtocolType {
    TCP,
    IB,
    PFC
};

// Use in controller
typedef Topology::Node topoNode;
typedef Topology::Link topoLink;

typedef pair<topoNode*, topoNode*> comm_pair;

// Describe the 'stage': that is th commpair, flows, routes
// of a task stage. So Ring has only 1 stage, butterfly has
// multiple stages.
// NOTE : DONT CHANGE THE STRUCT CASUALLY!
struct stageboard {
    bool newIterationFlag;
    // Commpairs of different stages
    vector<comm_pair> comm_pairs;
    // Links of different stages
    vector<topoLink*> used_links;
    // The key of the maps below are the GPUIdx
    // Flows of different stages
    map<int, vector<Flow>> flows_on_gpu;
    // The key of the maps below are the Module Id
    // Routes of different stages
    map<int, vector<RouteInfo>> routes_on_switch;
};

// The different stage means that the routes or flows has modified
// So the ring only has 1 stage, and HD have many stages
struct taskboard {
    Task the_task;
    int flows_count;
    omnetpp::simtime_t start_time;
    int working_workers;
    vector<topoNode*> used_gpus;
    vector<stageboard> stage_info;
};

struct ACLSingleEntry {
    unsigned int SourceAddressInt;
    unsigned int DestAddress;
    unsigned int NextHopAddress;
    int SourcePort;
    int interfaceID;
    int priority;
    int protocolID;
};

int log2(int n);

}  // namespace ai_sim

#endif  // CLUSTER_BASIC_BASIC_H_
